#!/usr/bin/env python
# -*- coding: utf-8 -*-

import abc
from typing import Dict, List, Any, Tuple, Optional, Union, Set
import os
from tree_sitter import Tree, Node
import logging

# 获取logger
logger = logging.getLogger(__name__)

class ASTToKGConverter(abc.ABC):
    """
    抽象语法树(AST)到知识图谱(KG)的转换器抽象接口
    可以为不同的编程语言实现具体子类
    """
    
    @abc.abstractmethod
    def initialize_parser(self) -> None:
        """
        初始化语法分析器
        每种语言需要实现自己的初始化逻辑
        """
        pass
    
    @abc.abstractmethod
    def parse_file(self, file_path: str) -> Optional[Tree]:
        """
        解析单个文件，生成语法树
        
        Args:
            file_path: 要解析的文件路径
            
        Returns:
            解析生成的语法树对象，解析失败时返回None
        """
        pass
    
    @abc.abstractmethod
    def parse_code(self, code: str) -> Optional[Tree]:
        """
        解析代码字符串，生成语法树
        
        Args:
            code: 要解析的代码字符串
            
        Returns:
            解析生成的语法树对象，解析失败时返回None
        """
        pass
    
    @abc.abstractmethod
    def parse_directory(self, directory_path: str, file_extension: str) -> Dict[str, Tree]:
        """
        解析目录中所有指定扩展名的文件
        
        Args:
            directory_path: 要解析的目录路径
            file_extension: 文件扩展名，例如'.java'
            
        Returns:
            文件路径到对应语法树的映射
        """
        pass
    
    @abc.abstractmethod
    def extract_entities(self, tree: Tree, file_path: str, source_code: str = None) -> Dict[str, List[Dict[str, Any]]]:
        """
        从语法树中提取实体（类、方法、变量等）
        
        Args:
            tree: 语法树对象
            file_path: 文件路径
            source_code: 源代码字符串（可选）
            
        Returns:
            按实体类型分组的实体列表
        """
        pass
    
    @abc.abstractmethod
    def extract_relationships(self, tree: Tree, file_path: str, entities: Dict[str, List[Dict[str, Any]]], source_code: str = None) -> List[Dict[str, Any]]:
        """
        从语法树中提取实体间的关系
        
        Args:
            tree: 语法树对象
            file_path: 文件路径
            entities: 已提取的实体
            source_code: 源代码字符串（可选）
            
        Returns:
            关系列表
        """
        pass
    
    @abc.abstractmethod
    def convert_to_graph(self, graph_connection, entities: Dict[str, List[Dict[str, Any]]], relationships: List[Dict[str, Any]]) -> bool:
        """
        将提取的实体和关系转换为图数据库中的节点和边
        
        Args:
            graph_connection: 图数据库连接对象
            entities: 按类型分组的实体列表
            relationships: 关系列表
            
        Returns:
            转换是否成功
        """
        pass
    
    def process_file(self, file_path: str, graph_connection) -> bool:
        """
        处理单个文件，提取实体和关系并存入图数据库
        
        Args:
            file_path: 文件路径
            graph_connection: 图数据库连接对象
            
        Returns:
            处理是否成功
        """
        try:
            # 读取源代码
            with open(file_path, 'r', encoding='utf-8') as f:
                source_code = f.read()
                
            tree = self.parse_file(file_path)
            if tree is None:
                return False
            
            entities = self.extract_entities(tree, file_path, source_code)
            relationships = self.extract_relationships(tree, file_path, entities, source_code)
            return self.convert_to_graph(graph_connection, entities, relationships)
        except Exception as e:
            logger.error(f"处理文件时发生错误: {str(e)}")
            return False
    
    def process_directory(self, directory_path: str, graph_connection, file_extension: str) -> Tuple[int, int]:
        """
        处理目录中所有指定扩展名的文件
        
        Args:
            directory_path: 目录路径
            graph_connection: 图数据库连接对象
            file_extension: 文件扩展名，例如'.java'
            
        Returns:
            (成功处理的文件数, 失败的文件数)
        """
        success_count = 0
        fail_count = 0
        
        for root, _, files in os.walk(directory_path):
            for file in files:
                if file.endswith(file_extension):
                    file_path = os.path.join(root, file)
                    if self.process_file(file_path, graph_connection):
                        success_count += 1
                    else:
                        fail_count += 1
        
        return success_count, fail_count
    
    @staticmethod
    def get_node_text(node: Node, code_bytes: bytes) -> str:
        """
        获取节点对应的源代码文本
        
        Args:
            node: 语法树节点
            code_bytes: 源代码的字节表示
            
        Returns:
            节点对应的文本
        """
        return code_bytes[node.start_byte:node.end_byte].decode('utf8')
    
    @staticmethod
    def find_all_nodes(root: Node, node_type: str) -> List[Node]:
        """
        在语法树中查找所有指定类型的节点
        
        Args:
            root: 根节点
            node_type: 节点类型
            
        Returns:
            节点列表
        """
        nodes = []
        
        def traverse(node):
            if node.type == node_type:
                nodes.append(node)
            for child in node.children:
                traverse(child)
        
        traverse(root)
        return nodes 